//
// DssFragmentation.cxx
//
// Created by TB on 11/18/11.
// Copyright 2011 BNL. All rights reserved.
//

#include <cmath>
#include <list>

#include <boost/foreach.hpp>
#include <boost/multi_array.hpp>

#include <TMath.h>
#include <TParticlePDG.h>
#include <TRandom3.h>
#include <TROOT.h>
#include <TVector3.h>
#include <TVirtualPad.h>

#include "tppmc/DssFragmentation.h"
#include "tppmc/KkpParton.h"

extern "C" {
   
   // Access to FORTRAN fDSS subroutine (see file fDSS.f).
   void fdss_(int* hadron, int* charge, int* order, double* z,
              double* scale, int* parton, double* result);
   
   // Access to FORTRAN FRAGINI COMMON block (see file fDSS.f).
   extern struct {
      double fini; 
   } fragini_;
   
} // extern "C"

namespace {
   enum {
      DUMMY = 0,
      SCALE = 1,
      PARTON = 2,
      ORDER = 3,
      HADRON = 4,
      CHARGE = 5
   };
   
   /*
    Wrapper around Fortran DSS subroutine (see fDSS.f).
    Compatible with ROOT TF1.
    x[0] is the z value at which to evaluate the function.
    Depends on six parameters, par[0] to par[5]:
    0: Fragmentation function type: DSS or DSS + Collins
    1: Q at which to evaluate (not Q^2)
    2: Parton type, following KKP convention (see fDSS.f)
    3: LO or NLO function
    4: Hadron species (pion, kaon, proton or charged hadrons)
    5: Hadron charge (+, - or average of +/-)
    */
   Double_t dssFragRootWrapper(Double_t* x, Double_t* par) {
      
      // Note implicit type conversions from par[...] to integers
      // for some parameters.
      double z = x[0];
      double scale = pow(par[SCALE], 2.); // Parameter stores Q, want Q^2
      int parton = par[PARTON];
      int order = par[ORDER]; // LO or NLO
      int hadron = par[HADRON]; // species
      int charge = par[CHARGE]; // charge
      double dss; // Value of the fragmentation function
      
      // Set the "first" flag in the DSS routine.
      static bool first = true;
      
      if(first) {
         fragini_.fini = 0;
         first = false; 
      } // if
      
      // If below charm/bottom limits then fragment as u/ubar
      if(((CQUARK == parton) or (CBAR == parton)) and (scale < (1.65 * 1.65))) {
         if(CQUARK == parton)
            parton = UQUARK;
         else
            parton = UBAR;
      } // if
      
      if(((BQUARK == parton) or (BBAR == parton)) and (scale < (4.85 * 4.85))) {
         /**
          \remark Shouldn't these assignments be the other way round
          for correct charge?
          */
         if(BQUARK == parton)
            parton = UQUARK;
         else
            parton = UBAR;
      } // if
      
      // Call the FORTRAN routine
      fdss_(&hadron,&charge,&order,&z,&scale,&parton,&dss);
      
      // The NLO set for protons (hadron==3) from gluons (parton==0)
      // returns small negative values for low q^2 (~1) and small z (~0.05).
      // Deal with this by setting the ff to zero. 
      
      if(dss < 0.) {
         if(not ((3 == hadron) and (0 == parton) and (z < 0.06))) {
            std::cerr << "DSS WARNING, D_h<0.0 (" << dss << ") for z = "
            << x[0] << " qs = " << par[1] << " parton = " << par[2] << " hadron = " 
            << par[4] << " order = " << par[3] << " charge = " << par[5] 
            << std::endl;
         } // if
         else {
            dss = 0.0; 
         } // else
      } // if
      
      return dss;
   }
}

namespace tppmc {

   //
   // class DssFragmentation
   //

   // Instantiate static fields
   
   int DssFragmentation::smNQBins(2);
   double DssFragmentation::smQMin(2.);
   double DssFragmentation::smQMax(102.);

   TF1* DssFragmentation::Function() {
      if(not mFunction) {
         int count(0);
         TString namebase("tppmc_dss_frag_function");
         TString name;
         do {
            name = namebase;
            name += count;
            ++count;
         }
         while(gROOT->FindObject(name));
         
         mFunction = new TF1J(name, dssFragRootWrapper,
                              mZLower, mZUpper, 6);
         // Set default parameters:
         mFunction->SetParameter(0, DSS); // This isn't actually used any more!
         mFunction->SetParameter(SCALE,  2.5);
         mFunction->SetParameter(PARTON, UQUARK);
         mFunction->SetParameter(ORDER,  mOrder);
         mFunction->SetParameter(HADRON, 1);
         mFunction->SetParameter(CHARGE, 1);
         mFunction->SetRange(mZLower, mZUpper);
      } // if
      return mFunction;
   }

   /*
    The c++ wrapper (DSSFragFunction) around the FORTRAN subroutine
    is itself called via a ROOT TF1 1D function.
    The TF1 takes 6 parameters (see DSSFragFunction for an explanation).
    */
   DssFragmentation::DssFragmentation(int order, double meanPt2)
   : mOrder(order)
   , mMeanPt2(meanPt2)
   , mFunction(NULL) {
   }

   DssFragmentation::~DssFragmentation() {
      if(mFunction) {
         delete mFunction;
         mFunction = NULL;
      } // if
   }

   void DssFragmentation::SetFunctionParameters(int htype){
      double species(-1.e100);
      double charge(-1.e100);
      
      // Set the hadron type code
      switch(htype) {
         case 211://PIPLUS:
         case -211://PIMINUS:
         case 111://PIZERO:
            species = 1.;
            break;
         case 311://KZERO:
         case -311://KZEROBAR:
         case 321://KPLUS:
         case -321://KMINUS:
            species = 2.;
            break;
         case 2212://PROTON:
         case -2212://PBAR:
         case 2112://NEUTRON:
         case -2112://NBAR:
            species = 3.;
            break;
      } // switch
      
      // Set the hadron charge
      switch(htype) {
         case 211://PIPLUS:
         case 321://KPLUS:
         case 2212://PROTON:
            charge = 1.;
            break;
         case -211://PIMINUS:
         case -321://KMINUS:
         case -2212://PBAR:
            charge = -1.;
            break;
         case 111://PIZERO:
         case 311://KZERO:
         case -311://KZEROBAR:
         case 2112://NEUTRON:
         case -2112://NBAR:
            charge = 0.;
            break;
      } // switch
      
      Function()->SetParameter(HADRON, species); 
      Function()->SetParameter(CHARGE, charge);
   }

   void DssFragmentation::Insert(int hadron) {
      // Add the parton/Q2 table for this hadron
      mTable.insert(std::make_pair(hadron,
                                   KeyedLists(smNQBins, smQMin, smQMax)));
      
      // Add the Q list for each parton in the new entry
      BOOST_FOREACH(int parton, Partons()) {
         HadronTable(hadron)->Insert(parton);
      } // for
   }

   bool DssFragmentation::Supports(const pythpart& parton) const {
      return std::find(Partons().begin(),
                       Partons().end(),
                       parton.type) not_eq Partons().end();
   }

   void DssFragmentation::PopulateKeyedLists() {
      // Save the parton type
      double saved_parton = Function()->GetParameter(PARTON);
      double saved_scale = Function()->GetParameter(SCALE);
      
      // Create a Q2/parton table for each hadron
      BOOST_FOREACH(int hadron, Hadrons()) {
         Insert(hadron);
      } // for
      
      /*
       For each (Q2 bin/hadron/parton), integrate the FF over
       the entire z range for which it is valid.
       Then normalise by the sum of integrals for all hadrons in
       that Q2 bin from that parton, so each entry is the probability 
       for that parton to fragment into that hadron type at that Q2.
       */
      BOOST_FOREACH(int parton, Partons()) { // Loop over parton types
         for(int qbin(0); qbin < smNQBins; qbin++) { // Loop over Q2 bins
            ComputeHadronProbabilities(parton, qbin);
         } // for
      } // BOOST_FOREACH
      
      Function()->SetParameter(PARTON, saved_parton);
      Function()->SetParameter(SCALE,  saved_scale);
   }

   void DssFragmentation::ComputeHadronProbabilities(int parton, int qbin) {
      // All hadron tables have the same bin width, so just
      // pick one to get the bin size.
      double QWidth = HadronTable(211/*PIPLUS*/)->BinSize();
      double binLowEdge = smQMin + QWidth * qbin;
      Function()->SetParameter(SCALE,  binLowEdge);
      Function()->SetParameter(PARTON, parton);
      
      double zLower(0.05), zUpper(1.);
      
      // DSS gives individual flavor/charge.
      // Calculate the integral of the FF for each hadron and accumulate
      // the total so we can normalise to 1 afterward.
      double sumIntegral(0.);
      BOOST_FOREACH(int hadron, Hadrons()) {
         // Set the function parameters for this Q2, parton and
         // compute the integral in this bin.
         SetFunctionParameters(hadron);
         double integral = Function()->Integral(zLower, zUpper);
         // Set the value for this bin.
         // Accumulate the integral so we can normalise later.
         HadronTable(hadron)->Set(parton, qbin, integral);
         sumIntegral += integral;
      } // for
      
      // Find total for all hadrons from this parton at this Q2
      // and normalise to give probability.
      if(fabs(sumIntegral) < 1.e-5) {
         // Oops, don't want to divide by zero...
         std::cerr << "Hadron probability integral for (parton, Q) = ("
         << parton << ", " << qbin << ") == 0!" << std::endl;
         SetErrorValues(parton, qbin);
      } // else
      else {
         // Normalise:
         BOOST_FOREACH(int hadron, Hadrons()) {
            double probability =
               HadronTable(hadron)->At(parton, qbin) / sumIntegral;
            HadronTable(hadron)->Set(parton, qbin, probability);
         } // for
      }  // if
   }

   void DssFragmentation::SetErrorValues(int q, int qbin) {
      std::cerr <<
      " ERROR in DssFragmentation - sumIntegral = 0, selecting pion"
      << std::endl;
      mTable[211].Set(q, qbin, 1. / 3.);
      mTable[-211].Set(q, qbin, 1. / 3.);
      mTable[111].Set(q, qbin, 1. / 3.);
      mTable[321].Set(q, qbin, 0.);
      mTable[-321].Set(q, qbin, 0.);
      mTable[311].Set(q, qbin, 0.);
      mTable[-311].Set(q, qbin, 0.);
      mTable[2212].Set(q, qbin, 0.);
      mTable[-2212].Set(q, qbin, 0.);
      mTable[2112].Set(q, qbin, 0.);
      mTable[-2112].Set(q, qbin, 0.);
   }

   const std::vector<int>& DssFragmentation::Hadrons() const {
      static std::vector<int> hadrons;
      if(hadrons.empty()) {
         int array[11] = {
            211, -211, 111, 311, -311, 321, -321,
            2212, -2212, 2112, -2112
         };
         hadrons.assign(array, array + 11);
      }
      return hadrons;
   }

   const std::vector<int>& DssFragmentation::Partons() const {
      static std::vector<int> partons;
      if(partons.empty()) {
         
         int array[11] = {
            GLUON,UQUARK,UBAR,DQUARK,DBAR,SQUARK,SBAR,
            CQUARK,CBAR,BQUARK,BBAR
         };
         partons.assign(array, array + 11);
      }
      return partons;
   }

   KeyedLists* DssFragmentation::HadronTable(int hadron) {
      
      if(mTable.empty()) PopulateKeyedLists();
      
      Table::iterator i = mTable.find(hadron);
      if(i not_eq mTable.end()) return &(i->second);
      
      return NULL;
   }

   double DssFragmentation::Evaluate(const pythpart& parton,
                                     const Hadron& hadron,
                                     double scale) {
      Function()->SetParameter(SCALE,  ::sqrt(scale));
      Function()->SetParameter(PARTON, kkp::fromPdg(parton.type));
      
      SetFunctionParameters(hadron.type);
      return Function()->Eval(hadron.z);
   }

   int DssFragmentation::ThrowHadronType(const pythpart& parton,
                                     double Q2) {
      if(INVALID_HADRON not_eq mHadron) {
         return mHadron;
      } // if
      std::vector<int> hadrons = Hadrons();
      
      double base(0.);
      
      int chosenHadron(INVALID_HADRON);
      const double select = gRandom->Uniform(1.);
      BOOST_FOREACH(int hadron, hadrons) {
         double val = HadronTable(hadron)->Interpolate(kkp::fromPdg(parton.type), sqrt(Q2));
         base += val;
         if(base > select) {
            chosenHadron = hadron;
            break;
         } // if
      } // BOOST_FOREACH
      // JPSI is the last
      if(INVALID_HADRON == chosenHadron) {
         std::cerr << "Threw invalid hadron type... using pi0" << std::endl;
         chosenHadron = 111;//PIZERO;
      } // if
      
      return chosenHadron;
   }

   std::vector<Hadron> DssFragmentation::Fragment(const pythpart& parton,
                                                  double Q2) {
      if(Q2 < 0.) {
         double pt2 = pow(parton.Pt(), 2.);
         if(pt2 < 2.) pt2 = 2.;
         Q2 = pt2;
      } // if
      
      int partontype = kkp::fromPdg(parton.Type());
      Function()->SetParameter(PARTON, partontype);
      Function()->SetParameter(SCALE,  ::sqrt(Q2));
      
      int h = ThrowHadronType(parton, Q2);
      Hadron hadron;
      hadron.type = h;
      Function()->SetRange(mZLower, mZUpper);
      hadron.z = Function()->GetRandom(mZLower, mZUpper);
      hadron.phigen = gRandom->Uniform(-TMath::Pi(), TMath::Pi());
      const double jt_sigma = sqrt(mMeanPt2/2.0);  
      double jt; 
      // Throw jt according to an exponential distribution such
      // that mean jt^2 is 0.2
      do {
         jt = jt_sigma*sqrt( -2.0*log(gRandom->Uniform(1.0)) ); 	  
      } while(jt>(3.0*jt_sigma));
      // Randomly generate phi.
      TVector3 mom(jt, 0., 0.);
      
      // TVector3::SetPhi rotates phi leaving magnitude and theta unchanged.
      mom.SetPhi(hadron.phigen);
      hadron.px = mom.Px();
      hadron.py = mom.Py();
      
      hadron.mParton = parton;
      hadron.mFrag = Evaluate(parton, hadron, Q2);
      std::vector<Hadron> products;
      products.push_back(hadron);
      return products;
   }
} // namespace tppmc
